use godot::{
    classes::{notify::ContainerNotification, *},
    prelude::*,
};

#[derive(GodotClass)]
#[class(init, tool, base=Container)]
pub struct AspectRatioGridContainer {
    base: Base<Container>,
    #[export(range=(0.001, 10.0, 0.0001, or_greater))]
    #[var(get, set=set_ratio)]
    #[init(val = 1.0)]
    ratio: f32,
    #[export]
    #[var(get, set=set_spacing)]
    spacing: Vector2,
    #[export]
    #[var(get, set=set_columns)]
    #[init(val = 1)]
    columns: i32,
    #[export(enum = (Horizontal = 0, Vertical = 1))]
    #[var(get, set=set_layout_direction)]
    #[init(val = 0)]
    layout_direction: i32,
}

#[inline]
fn to_axis(i: i32) -> Vector2Axis {
    unsafe { std::mem::transmute(i) }
}

#[godot_api]
impl IContainer for AspectRatioGridContainer {
    fn on_notification(&mut self, what: ContainerNotification) {
        if what == ContainerNotification::SORT_CHILDREN {
            let main_axis = to_axis(self.layout_direction);
            let sub_axis = to_axis((self.layout_direction + 1) % 2);

            let cell_stride = (self.base().get_rect().size[main_axis]
                + self.spacing[main_axis])
                / self.columns as f32;

            let mut child_size = Vector2::splat(cell_stride - self.spacing[main_axis]);
            if let Vector2Axis::X = sub_axis {
                child_size[sub_axis] *= self.ratio;
            } else {
                child_size[sub_axis] /= self.ratio;
            }
            
            let cell_size = child_size + self.spacing;

            for idx in 0..self.base().get_child_count() {
                if let Some(Ok(child)) = self.base().get_child(idx).map(Gd::try_cast::<Control>) {
                    let mut cell = Vector2::ZERO;
                    cell[main_axis] = (idx % self.columns) as _;
                    cell[sub_axis] = (idx / self.columns) as _;

                    let pos = cell_size * cell;

                    self.base_mut()
                        .fit_child_in_rect(&child, Rect2::new(pos, child_size))
                }
            }

            self.base_mut().update_minimum_size();
        }
    }

    fn get_minimum_size(&self) -> Vector2 {
        let rows = self.base().get_child_count() as f32 / self.columns as f32;
        let sub_axis = to_axis((self.layout_direction + 1) % 2);
        let mut ret = Vector2::ZERO;
        ret[sub_axis] = rows.ceil() * self.get_cell_size(sub_axis) - self.spacing[sub_axis];
        ret
    }
}

impl AspectRatioGridContainer {
    pub fn get_cell_size(&self, axis: Vector2Axis) -> f32 {
        if axis == to_axis(self.layout_direction) {
            (self.base().get_rect().size[axis] + self.spacing[axis]) / self.columns as f32
        } else {
            let main_axis = to_axis(self.layout_direction);
            let main_size = self.get_cell_size(main_axis);

            if axis == Vector2Axis::X {
                (main_size - self.spacing[main_axis]) * self.ratio + self.spacing[axis]
            } else {
                (main_size - self.spacing[main_axis]) / self.ratio + self.spacing[axis]
            }
        }
    }
}

#[godot_api]
impl AspectRatioGridContainer {
    #[func]
    pub fn set_ratio(&mut self, ratio: f32) {
        if ratio == self.ratio {
            return;
        }
        self.ratio = ratio;

        self.base_mut().queue_sort();
        self.base_mut().update_minimum_size()
    }

    #[func]
    pub fn set_columns(&mut self, columns: i32) {
        if columns == self.columns {
            return;
        }
        self.columns = columns;

        self.base_mut().queue_sort();
        self.base_mut().update_minimum_size()
    }

    #[func]
    pub fn set_spacing(&mut self, spacing: Vector2) {
        if spacing == self.spacing {
            return;
        }

        self.spacing = spacing;

        self.base_mut().queue_sort();
        self.base_mut().update_minimum_size()
    }

    #[func]
    pub fn set_layout_direction(&mut self, alignment: i32) {
        if alignment == self.layout_direction {
            return;
        }

        self.layout_direction = alignment;

        self.base_mut().queue_sort();
        self.base_mut().update_minimum_size()
    }
}
